#include <rtthread.h>
#include <board.h>
#include <fal.h>
#include <dfs.h>
#include <dfs_file.h>
#include <dfs_private.h>

#include "bl_if.h"
#include "bl_crc32.h"
#include "tiny_aes.h"
#include "quicklz.h"
#include "fastlz.h"

#ifdef TBOOT_USING_USB
extern rt_bool_t rt_has_mass_disk(void);

rt_bool_t updateFromUSBDisk(void)
{
    int retry = 5;
    struct stat stat;
    struct dfs_fd fd_src;

    UpdateProgress uPState = TB_READ_MORE;
    rt_int32_t numb = 0, rd_cnt = 0, wr_cnt = 0, wr_off = 0;
    bl_header_info_t header = RT_NULL;

    rt_uint8_t *ctx_buf = RT_NULL, *buf = RT_NULL;
    rt_uint32_t crc = 0xFFFFFFFF;
    rt_uint32_t compr_blk_sz = 0;
    rt_uint32_t compr_off = 0;
    rt_size_t decompr_sz = 0;
    rt_size_t blk_sz_off = 0, cpy_len = 0, decrypt_off = 0;
    rt_uint8_t *compr_buf = RT_NULL, *decompr_buf = RT_NULL;
    rt_uint8_t *decrypt_buf = RT_NULL;
#if defined TBOOT_USING_QUICKLZ
    qlz_state_decompress *qlz_decompr_state = RT_NULL;
#endif
#ifdef TBOOT_USING_AES256
	tiny_aes_context *aes_ctx = RT_NULL;
	rt_uint8_t *aes_iv = RT_NULL;
#endif

    // waitfor usb disk
    retry = 5;
    while (retry > 0) {
        if (rt_has_mass_disk() == RT_TRUE) {
            break;
        }
        rt_thread_mdelay(10);
        retry--;
    }
    if (retry == 0) return RT_FALSE;

    // check file
    retry = 5;
    while (retry > 0) {
        if (dfs_file_stat(BOOT_APP_FILE_PATH, &stat) == 0) {
            break;
        }
        rt_thread_mdelay(10);
        retry--;
    }
    if (retry == 0) return RT_FALSE;

    /* verify */
    // open new application file
    if ((dfs_file_open(&fd_src, BOOT_APP_FILE_PATH, O_RDONLY) != 0)) {
        return RT_FALSE;
    }
    // read file header
    header = rt_malloc(sizeof(bl_header_info));
    if (header == RT_NULL) return RT_FALSE;
    numb = dfs_file_read(&fd_src, header, sizeof(bl_header_info));
    if ( numb != sizeof(bl_header_info)) {
        rt_free(header);
        return RT_FALSE;
    }
    // check header crc sum
    crc = bl_crc32sum((rt_uint8_t *)header, sizeof(bl_header_info)-4);
    if (header->header_crc != crc) {
        rt_free(header);
        dfs_file_close(&fd_src);
        return RT_FALSE;
    }
    // check bootloader version
    if ((header->tbl_ver >> 24) > MAJOR_VERSION) {
        rt_free(header);
        dfs_file_close(&fd_src);
        return RT_FALSE;
    }

    // check magic numb
    if (header->magic_numb != BOOT_MAGIC_NUMB) {
        rt_free(header);
        dfs_file_close(&fd_src);
        return RT_FALSE;
    }
    // check model
    if (rt_strncmp(header->model, BOOT_INSTRUMENT_MODEL, sizeof(header->model)) != 0) {
        rt_free(header);
        dfs_file_close(&fd_src);
        return RT_FALSE;
    }
    // check context size
    if (stat.st_size != header->ctx_size + sizeof(bl_header_info)) {
        rt_free(header);
        dfs_file_close(&fd_src);
        return RT_FALSE;
    }

    // read all context
    ctx_buf = rt_malloc(header->ctx_size);
    if (ctx_buf == RT_NULL) {
        rt_free(header);
        dfs_file_close(&fd_src);
        return RT_FALSE;
    }
    numb = dfs_file_read(&fd_src, ctx_buf, header->ctx_size);
    dfs_file_close(&fd_src);
    if (numb != header->ctx_size) {
        goto _err_exit;
    }
    // check context crc
    crc = 0xFFFFFFFF;
    crc = bl_crc32(crc, ctx_buf, numb);
    crc ^= 0xFFFFFFFF;
    if (header->ctx_crc != crc) {
        goto _err_exit;
    }

    // malloc read buffer
    buf = rt_malloc(TBOOT_READ_BUFF_SZ);
    if (buf == RT_NULL) {
        goto _err_exit;
    }

    // decryption
    if ((header->flag & BOOT_ENCRYP_MASK) != BOOT_ENCRYP_NONE) {
#ifdef TBOOT_USING_AES256
        if ((header->flag & BOOT_ENCRYP_MASK) == BOOT_ENCRYP_AES256) {
            aes_ctx = rt_malloc(sizeof(tiny_aes_context));
            if (aes_ctx == RT_NULL) {
                goto _err_exit;
            }
            aes_iv = rt_malloc(16);
            if (aes_iv == RT_NULL) {
                goto _err_exit;
            }

            rt_memcpy(aes_iv, BOOT_ENCRYP_AES_KEY, 16);
            tiny_aes_setkey_dec(aes_ctx, (rt_uint8_t *)BOOT_ENCRYP_AES_KEY, 256);
            decrypt_buf = rt_malloc(TBOOT_READ_BUFF_SZ);
        } else
#endif
        {
            goto _err_exit;
        }
    } else {
        decrypt_buf = buf;
    }
    
    if ((header->flag & BOOT_COMPRE_MASK) != BOOT_COMPRE_NONE) {
#if defined TBOOT_USING_FASTLZ
        if ((header->flag & BOOT_COMPRE_MASK) == BOOT_COMPRE_FASTLZ) {
            compr_buf = rt_malloc(TBOOT_COMP_BUFF_SZ);
            decompr_buf = rt_malloc(TBOOT_COMP_BUFF_SZ);
        } else 
#elif defined TBOOT_USING_QUICKLZ
        if ((header->flag & BOOT_COMPRE_MASK) == BOOT_COMPRE_QUICKLZ) {
            qlz_decompr_state = rt_malloc(sizeof(qlz_state_decompress));
            rt_memset(qlz_decompr_state, 0x0, sizeof(qlz_state_decompress));
            compr_buf = rt_malloc(TBOOT_COMP_BUFF_SZ);
            decompr_buf = rt_malloc(TBOOT_COMP_BUFF_SZ);
        } else 
#endif
        {
            goto _err_exit;
        }
    } else {
        decompr_buf = rt_malloc(TBOOT_COMP_BUFF_SZ);
    }

    // erase on chip flash
    dstPartition = fal_partition_find(header->parti_name);
	if (dstPartition == RT_NULL) {
		goto _err_exit;
	}
    if (fal_partition_erase(dstPartition, 0, dstPartition->len) < 0) {
        goto _err_exit;
    }

    rt_thread_yield();
//    crc = 0xFFFFFFFF;
    while (uPState != TB_FIN) {
    switch(uPState){
    case TB_READ_MORE:
#if 0
        numb = dfs_file_read(&fd_src, buf, TBOOT_READ_BUFF_SZ);
        if (numb <= 0) {
            uPState = TB_FIN;
            break;
        } else if (numb != TBOOT_READ_BUFF_SZ) {
//            printf("%d", numb);
        }
#else
        if (rd_cnt == header->ctx_size) {
            uPState = TB_FIN;
            break;
        }
        if (rd_cnt + TBOOT_READ_BUFF_SZ < header->ctx_size) {
            numb = TBOOT_READ_BUFF_SZ;
        } else {
            numb = header->ctx_size - rd_cnt;
        }
        memcpy(buf, &ctx_buf[rd_cnt], numb);
#endif
        rd_cnt += numb;
//        crc = bl_crc32(crc, buf, numb);
    case TB_DECRPT:
        // decryption
#ifdef TBOOT_USING_AES256
        if ((header->flag & BOOT_ENCRYP_MASK) != BOOT_ENCRYP_NONE) {
            tiny_aes_crypt_cbc(aes_ctx, AES_DECRYPT, numb, aes_iv, buf, &decrypt_buf[decrypt_off]);
            decrypt_off += numb;
        } else 
#endif
        {
            decrypt_off += numb;
        }
    case TB_DECOMPR:
        // decompress
        if ((header->flag & BOOT_COMPRE_MASK) != BOOT_COMPRE_NONE) {
            if (decrypt_off < TBOOT_COMP_HEAD_SZ) {
                uPState = TB_READ_MORE;
                break;
            }
            if (compr_off == 0) {
                compr_blk_sz = decrypt_buf[0] << 24 | decrypt_buf[1] << 16 | decrypt_buf[2] << 8 | decrypt_buf[3];
                if (compr_blk_sz > dstPartition->len) {        // error!!!
                    uPState = TB_FIN;
                    break;
                }
                blk_sz_off = TBOOT_COMP_HEAD_SZ;
            } else {
                blk_sz_off = 0;
            }
            if ((compr_blk_sz - compr_off) < (decrypt_off - blk_sz_off)) {
                cpy_len = compr_blk_sz - compr_off;
                rt_memcpy(&compr_buf[compr_off], decrypt_buf+blk_sz_off, cpy_len);
                decrypt_off = decrypt_off - blk_sz_off - cpy_len;
                rt_memmove(decrypt_buf, decrypt_buf+blk_sz_off+cpy_len, decrypt_off);
            } else {
                cpy_len = decrypt_off - blk_sz_off;
                rt_memcpy(&compr_buf[compr_off], decrypt_buf+blk_sz_off, cpy_len);
                decrypt_off = 0;
            }
            compr_off += cpy_len;

            if (compr_off < compr_blk_sz) { // read more block
                uPState = TB_READ_MORE;
                break;
            }

#if defined TBOOT_USING_FASTLZ
            if ((header->flag & BOOT_COMPRE_MASK) == BOOT_COMPRE_FASTLZ) {
                decompr_sz = fastlz_decompress((const void *)compr_buf, compr_off, decompr_buf, TBOOT_READ_BUFF_SZ);
            }
#elif defined TBOOT_USING_QUICKLZ
            if ((header->flag & BOOT_COMPRE_MASK) == BOOT_COMPRE_QUICKLZ) {
                decompr_sz = qlz_decompress((const char *)compr_buf, decompr_buf, qlz_decompr_state);
            }
#endif
            if (decompr_sz <= 0) {        // error!!!
                uPState = TB_FIN;
                break;
            }
        } else {
            if (wr_off + decrypt_off > header->raw_size) {
                decrypt_off = header->raw_size - wr_off;
            }
            rt_memcpy(decompr_buf, decrypt_buf, decrypt_off);
            decompr_sz = decrypt_off;
            decrypt_off = 0;
        }
    case TB_WRITE:
        if (wr_off + decompr_sz > dstPartition->len) {        // error!!!
            uPState = TB_FIN;
            break;
        }
        wr_cnt = fal_partition_write(dstPartition, wr_off, decompr_buf, decompr_sz);
        if (wr_cnt != decompr_sz) {        // error!!!
            uPState = TB_FIN;
            break;
        }
        wr_off += decompr_sz;
        rt_thread_yield();
    case TB_PRE_NEXT:
        if ((header->flag & BOOT_COMPRE_MASK) != BOOT_COMPRE_NONE) {
            compr_off = 0;
            if (decrypt_off > 0) {
                uPState = TB_DECOMPR;
                break;
            }
        }
        uPState = TB_READ_MORE;
    break;
    case TB_FIN:
    break;
    default:
    break;
    }
    }
//    crc ^= 0xFFFFFFFF;

_err_exit:
    if (buf) rt_free(buf);
    if (ctx_buf) rt_free(ctx_buf);
#ifdef TBOOT_USING_AES256
    if (aes_iv) rt_free(aes_iv);
    if (aes_ctx) rt_free(aes_ctx);
#endif
#if defined TBOOT_USING_QUICKLZ
    if (qlz_decompr_state) rt_free(qlz_decompr_state);
#endif
    if (decrypt_buf) rt_free(decrypt_buf);
    if (compr_buf) rt_free(compr_buf);
    if (decompr_buf) rt_free(decompr_buf);

    if (header->raw_size != wr_off) {
        rt_free(header);
        return RT_FALSE;
    }
    rt_free(header);

    return RT_TRUE;
}

#endif /* TBOOT_USING_USB */
